import * as tf from '@tensorflow/tfjs'
import * as tfvis from '@tensorflow/tfjs-vis'

window.onload = async () => {
  const xs = [1, 2, 3, 4]
  const ys = [1, 3, 5, 7]

  tfvis.render.scatterplot(
    { name: '线性回归训练集', tab: 'charts' },
    { values: xs.map((x, i) => ({ x, y: ys[i] })), series: ['trend'] },
    { xAxisDomain: [0, 5], yAxisDomain: [0, 8] }
  )

  // 初始化模型
  const model = tf.sequential()
  // 添加层
  model.add(tf.layers.dense({ units: 1, inputShape: [1] }))
  // 损失函数，优化器
  model.compile({
    loss: tf.losses.meanSquaredError, // 损失函数
    optimizer: tf.train.sgd(0.1), // 优化器,学习率
  })

  const inputs = tf.tensor(xs)
  const labels = tf.tensor(ys)
  await model.fit(inputs, labels, {
    batchSize: 4,
    epochs: 100,
    callbacks: tfvis.show.fitCallbacks(
      {
        name: '训练过程',
      },
      ['loss']
    ),
  })

  let guessX = 5
  const output = model.predict(tf.tensor([guessX]))

  output.print()
  console.log(`如果x为 ${guessX}, 输出有为 ${output.dataSync()[0]}`)
}
